// arch/riscv/kernel/proc.c
#include "proc.h"
#include "mm.h"
#include "rand.h"
#include "printk.h"
#include "defs.h"
#include "vm.h"
#include "elf.h"
#include "string.h"
// arch/riscv/kernel/proc.c

extern void __dummy();
extern unsigned long *swapper_pg_dir;
extern uint64_t _sramdisk;
extern uint64_t _eramdisk;

struct task_struct *idle;           // idle process
struct task_struct *current;        // 指向当前运行线程的 `task_struct`
struct task_struct *task[NR_TASKS]; // 线程数组, 所有的线程都保存在此

/**
 * new content for unit test of 2023 OS lab2
 */
extern uint64 task_test_priority[]; // test_init 后, 用于初始化 task[i].priority 的数组
extern uint64 task_test_counter[];  // test_init 后, 用于初始化 task[i].counter  的数组

static uint64_t load_elf_program(struct task_struct *task)
{
    Elf64_Ehdr *ehdr = (Elf64_Ehdr *)(&_sramdisk);

    uint64_t phdr_start = (uint64_t)ehdr + ehdr->e_phoff;
    int phdr_cnt = ehdr->e_phnum;

    Elf64_Phdr *phdr;
    int load_phdr_cnt = 0;
    for (int i = 0; i < phdr_cnt; i++)
    {   
        phdr = (Elf64_Phdr *)(phdr_start + sizeof(Elf64_Phdr) * i);
        if (phdr->p_type == PT_LOAD)
        {
            // 通过p_memsz来计算拷贝该Segment所需要的页数量，并用过alloc_page申请页
            uint64 s_page = alloc_pages((phdr->p_memsz - 1) / PGSIZE + 1);
            // 从elf文件中读取Segment内容，拷贝到之前申请的页上
            // ((phdr->p_vaddr) - PGROUNDDOWN(phdr->p_vaddr))为该Segment的内容相对于page起始位置的offset,从而保持其相对位置不变
            // ((uint64)(&uapp_start) + phdr->p_offset)为该Segment的内容在uaqq中的位置
            memcpy((void *)(s_page + ((phdr->p_vaddr) - PGROUNDDOWN(phdr->p_vaddr))), (void *)((uint64)(&_sramdisk) + phdr->p_offset), phdr->p_memsz);
            // 映射页表
            // VA2PA(s_page)所表示的物理地址映射至PGROUNDDOWN(phdr->p_vaddr)所对应的虚拟地址
            // (phdr->p_flags << 1 | 17)将页表项的Valid设置为1，并将该页表设置为U-Mode
            create_mapping((uint64 *)PA2VA((uint64_t)task->pgd), PGROUNDDOWN(phdr->p_vaddr), VA2PA(s_page), phdr->p_memsz, (phdr->p_flags << 1 | 17));
        }
    }

    // 拷贝完成后，为用户进程在其内存的尾部添加新的页，作为用户进程的栈，同样通过页表映射至用户进程的虚拟内存上
    create_mapping((uint64 *)PA2VA((uint64)task->pgd), USER_END - PGSIZE, VA2PA(alloc_page()), PGSIZE, 23);

    // sepc设置为用户程序的第一条指令被存储的虚拟地址
    task->thread.sepc = ehdr->e_entry;
    // 设置 sstatus 的 SPP = 0, SPIE = 1, SUM = 1
    task->thread.sstatus = 0 | 1 << 18 | 1 << 5;
    // 设置 sscratch 为 尾部的USER_END
    task->thread.sscratch = USER_END;
}

void task_init()
{
    //test_init(NR_TASKS);
    printk("tash_init start\n");
    idle = (struct task_struct *)kalloc(); // 调用 kalloc() 为 idle 分配一个物理页
    idle->state = TASK_RUNNING;            // 设置 state 为 TASK_RUNNING;
    idle->counter = idle->priority = 0;    // idle 不参与调度，将其 counter / priority 设置为 0
    idle->pid = 0;                         // 设置 idle 的 pid 为 0
    current = idle;                        // 将 current 和 task[0] 指向 idle
    task[0] = idle;

    //设置页表、sscratch
    idle->pgd = swapper_pg_dir;
    idle->thread.sscratch = 0;

    int i;
    // 通过指令获取__dummy的地址
    uint64 dummy_address;
    __asm__ volatile(
        "la %0, __dummy"
        : "=r"(dummy_address)
        :
        : "memory");
    // 设置每一个线程，过程与设置idle类似
    for (i = 1; i < NR_TASKS; i++)
    {
        task[i] = (struct task_struct *)kalloc();
        task[i]->pid = i;
        task[i]->state = TASK_RUNNING;
        //task[i]->counter = task_test_counter[i];
        task[i]->counter = 0;
        task[i]->priority = rand() % PRIORITY_MAX + 1;
        //task[i]->priority = task_test_priority[i];
        printk("[proc_init] %d  counter: %ld priority: %ld\n", i, task[i]->counter, task[i]->priority);
        task[i]->thread.ra = dummy_address;            // ra 设置为 __dummy 的地址
        task[i]->thread.sp = (uint64)task[i] + PGSIZE; // sp 设置为该线程申请的物理页的高地址，即低地址+物理页大小
        
        //申请新的页作U态页表，并将S态页表复制进来
        pagetable_t pgt = (pagetable_t)alloc_page();
        memcpy((void *)pgt, (void *)(&swapper_pg_dir), PGSIZE);
        task[i]->pgd = (pagetable_t)VA2PA((uint64_t)pgt); 
        
        /* 不使用load_elf_programm
        // 通过alloc_page申请页
        uint64 s_page = alloc_page ((uapp_end - uapp_start - 1) / PGSIZE + 1);
        // 从elf文件中读取Segment内容，拷贝到之前申请的页上
        memcpy((void *)(s_page), (void *)((uint64_t)(uapp_start)), uapp_end - uapp_start);
        // 映射页表
        create_mapping((uint64 *)PA2VA((uint64_t)task->pgd), PGROUNDDOWN(USER_START), VA2PA(s_page), uapp_end - uapp_start, 31);
        // 映射完成后，为用户进程在其内存的尾部添加新的页，作为用户进程的栈，同样通过页表映射至用户进程的虚拟内存上
        create_mapping(PA2VA((uint64)task->pgd), USER_END - PGSIZE, VA2PA(alloc_page()), PGSIZE, 23);
        // sepc设置为用户程序的第一条指令被存储的虚拟地址        
        task[i]->thread.sepc = USER_START;
        // 设置 sstatus 的 SPP = 0, SPIE = 1, SUM = 1
        task[i]->thread.sstatus = (1 << 5) | (1 << 18);
        // 设置 sscratch 为 尾部的USER_END
        task[i]->thread.sscratch = USER_END;
        */
        load_elf_program(task[i]);

        task[i]->files = file_init();
    }
    printk("...proc_init done!\n");
}

// arch/riscv/kernel/proc.c
void dummy()
{
    //schedule_test();
    // printk("AAA s9:%lld s10:%lld\n", idle->thread.s[9], idle->thread.s[10]);
    uint64 MOD = 1000000007;
    uint64 auto_inc_local_var = 0;
    int last_counter = -1;
    while (1)
    {
        if ((last_counter == -1 || current->counter != last_counter) && current->counter > 0)
        {
            if (current->counter == 1)
            {
                --(current->counter); // forced the counter to be zero if this thread is going to be scheduled
            }                         // in case that the new counter is also 1, leading the information not printed.
            last_counter = current->counter;
            auto_inc_local_var = (auto_inc_local_var + 1) % MOD;
            printk("[PID = %d] is running. auto_inc_local_var = %d, thread space begin at 0x%llx\n", current->pid, auto_inc_local_var, current);
        }
    }
}

extern void __switch_to(struct task_struct *prev, struct task_struct *next);

void switch_to(struct task_struct *next)
{
    // printk("switch_to %lld, counter = %ld\n", next->pid, next->counter);
    if (next->pid == current->pid) // 下一个执行的线程与当前的线程是同一个线程, 无需做任何处理
        return;
    // 下一个执行的线程与当前的线程不是同一个线程，进行进一步操作
    struct task_struct *prev = current;
    current = next; // 更新current指针，指向下一个进程
    __switch_to(prev, next);
}

void do_timer(void)
{
    // printk("[PID = %d] is do_timer. counter = %d   ", current->pid, current->counter);
    if (current->pid == 0) // 如果当前线程是 idle 线程，直接进行调度
        schedule();
    else
    {
        if (current->counter > 0) // 对当前线程的运行剩余时间减1
            current->counter--;
        if (current->counter == 0) // 若剩余时间等于0 则进行调度
            schedule();
    }
}
//短作业优先
void schedule(void)
{
    uint64 minn = 20;
    int _id = -1;
    int i;
    for (i = 1; i < NR_TASKS; i++)
    {
        if (task[i]->counter > 0 && task[i]->counter < minn) // 优先选择非0的剩余时间最少的线程
            _id = i, minn = task[i]->counter;
    }
    if (_id != -1)
    {
        switch_to(task[_id]); // 找到下一个线程，进行切换
        return;
    }
    else
    {
        // 所有线程的剩余时间均耗尽，通过调用rand()函数为每个线程重新随机赋予时间
        minn = 20;
        for (i = 1; i < NR_TASKS; i++)
        {
            task[i]->counter = rand();
            //task[i]->counter = 1;
            if (task[i]->counter > 0 && task[i]->counter < minn)
                _id = i, minn = task[i]->counter;
        }
        switch_to(task[_id]);
        return;
    }
}
//优先级优先
void schedule_(void)
{
    uint64 minn = 20;
    uint64 maxx = 0;
    int _id = -1;
    int i;
    for (i = 1; i < NR_TASKS; i++)
    {
        if (task[i]->counter > 0 && task[i]->counter >= maxx)
            _id = i, maxx = task[i]->counter;
    }
    if (_id != -1)
    {
        switch_to(task[_id]);
        return;
    }
    else
    {
        maxx = 0;
        for (i = 1; i < NR_TASKS; i++)
        {
            task[i]->counter = task[i]->priority;
            if (task[i]->counter > 0 && task[i]->counter > maxx)
                _id = i, maxx = task[i]->counter;
        }
        switch_to(task[_id]);
        return;
    }
}